# -*- coding: utf-8 -*-

"""
Augmentations used in object detection.

Input image format is: [height, row, channel], dtype is np.uint8
Input annotations format is: [[class_id, c_x, c_y, w, h], ..., [class_id, c_x, c_y, w, h]]
"""


import numpy as np
import tensorlayer as tl


def _obj_box_vertical_flip(im, coords=None, is_rescale=False, is_center=False, is_random=False):
    if coords is None:
        coords = []

    def _flip(im, coords):
        im = tl.prepro.flip_axis(im, axis=0, is_random=False)
        coords_new = list()

        for coord in coords:

            if len(coord) != 4:
                raise AssertionError("coordinate should be 4 values : [x, y, w, h]")

            if is_rescale:
                if is_center:
                    # x_center' = 1 - x
                    y = 1. - coord[1]
                else:
                    # x_center' = 1 - x - w
                    y = 1. - coord[1] - coord[3]
            else:
                if is_center:
                    # x' = im.width - x
                    y = im.shape[0] - coord[1]
                else:
                    # x' = im.width - x - w
                    y = im.shape[0] - coord[1] - coord[3]
            coords_new.append([coord[0], y, coord[2], coord[3]])
        return im, coords_new

    if is_random:
        factor = np.random.uniform(-1, 1)
        if factor > 0:
            return _flip(im, coords)
        else:
            return im, coords
    else:
        return _flip(im, coords)


def random_horizontal_flip(image, annotations):
    """Horizontal flip an image randomly.

    Parameters
    ----------
    image : numpy.array
        An image with dimension of [row, col, channel] (default).
    annotations : list of list of 5 float elements
        Annotations [[class_id, c_x, c_y, w, h], [class_id, c_x, c_y, w, h], ...].

    Returns
    -------
        numpy.array
            A processed image
        list of list of 5 float elements
            A list of new annotations.
    """
    image, coords = tl.prepro.obj_box_horizontal_flip(
        im=image,
        coords=[e[1:] for e in annotations],
        is_rescale=True,
        is_center=True,
        is_random=True)
    annotations = [[a[0]] + b for a, b in zip(annotations, coords)]
    return image, annotations


def random_vertical_flip(image, annotations):
    """Vertical flip an image randomly.

    Parameters
    ----------
    image : numpy.array
        An image with dimension of [row, col, channel] (default).
    annotations : list of list of 5 float elements
        Annotations [[class_id, c_x, c_y, w, h], [class_id, c_x, c_y, w, h], ...].

    Returns
    -------
        numpy.array
            A processed image
        list of list of 5 float elements
            A list of new annotations.
    """
    image, coords = _obj_box_vertical_flip(
        im=image,
        coords=[e[1:] for e in annotations],
        is_rescale=True,
        is_center=True,
        is_random=True)
    annotations = [[a[0]] + b for a, b in zip(annotations, coords)]
    return image, annotations


def shift(image, annotations, horizontal_shift_ratio=0.1, vertical_shift_ratio=0.1):
    """Shift an image.

    Parameters
    ----------
    image : numpy.array
        An image with dimension of [row, col, channel] (default).
    annotations : list of list of 5 elements
        Annotations [[class_id, c_x, c_y, w, h], [class_id, c_x, c_y, w, h], ...].
    horizontal_shift_ratio : float
        Horizontal shift ratio, positive for right shift, negative for left shift.
    vertical_shift_ratio : float
        Vertical shift ratio, positive for bottom shift, negative for top shift.

    Returns
    -------
        numpy.array
            A processed image
        list of list of 5 float elements
            A list of new annotations.
    """
    image, classes, coords = tl.prepro.obj_box_shift(
        im=image,
        classes=[e[0] for e in annotations],
        coords=[e[1:] for e in annotations],
        wrg=horizontal_shift_ratio,
        hrg=vertical_shift_ratio,
        is_rescale=True,
        is_center=True,
        is_random=False)
    annotations = [[a] + b for a, b in zip(classes, coords)]
    return image, annotations


def zoom(image, annotations, width_zoom_ratio, height_zoom_ratio):
    """Zoom an image.

    Parameters
    ----------
    image : numpy.array
        An image with dimension of [row, col, channel] (default).
    annotations : list of list of 5 elements
        Annotations [[class_id, c_x, c_y, w, h], [class_id, c_x, c_y, w, h], ...].
    width_zoom_ratio : float
        Width zoom ratio, less than 1 for zoom in, larger than 1 for zoom out.
    height_zoom_ratio : float
        Height zoom ratio, less than 1 for zoom in, larger than 1 for zoom out.

    Returns
    -------
        numpy.array
            A processed image
        list of list of 5 float elements
            A list of new annotations.
    """
    image, classes, coords = tl.prepro.obj_box_zoom(
        im=image,
        classes=[e[0] for e in annotations],
        coords=[e[1:] for e in annotations],
        zoom_range=(height_zoom_ratio, width_zoom_ratio),
        is_rescale=True,
        is_center=True,
        is_random=False)
    annotations = [[a] + b for a, b in zip(classes, coords)]
    return image, annotations


def resize(image, annotations, dst_size):
    """Resize image.

    Parameters
    ----------
    image : numpy.array
        An image with dimension of [row, col, channel] (default).
    annotations : list of list of 5 elements
        Annotations [[class_id, c_x, c_y, w, h], [class_id, c_x, c_y, w, h], ...].
    dst_size : tuple of 2 int
        For height and width.

    Returns
    -------
        numpy.array
            A processed image
        list of list of 5 float elements
            A list of new annotations.
    """
    image, coords = tl.prepro.obj_box_imresize(
        im=image,
        coords=[e[1:] for e in annotations],
        size=dst_size,
        is_rescale=True)
    annotations = [[a[0]] + b for a, b in zip(annotations, coords)]
    return image, annotations


def central_crop(image, annotations, dst_size):
    """Central crop from an image.

    Parameters
    ----------
    image : numpy.array
        An image with dimension of [row, col, channel] (default).
    annotations : list of list of 5 elements
        Annotations [[class_id, c_x, c_y, w, h], [class_id, c_x, c_y, w, h], ...].
    dst_size : tuple of 2 int
        For cropped height and width.

    Returns
    -------
        numpy.array
            A processed image
        list of list of 5 float elements
            A list of new annotations.
    """
    hrg, wrg = dst_size
    image, classes, coords = tl.prepro.obj_box_crop(
        im=image,
        classes=[e[0] for e in annotations],
        coords=[e[1:] for e in annotations],
        wrg=wrg,
        hrg=hrg,
        is_rescale=True,
        is_center=True,
        is_random=False)
    annotations = [[a] + b for a, b in zip(classes, coords)]
    return image, annotations


def random_crop(image, annotations, dst_size):
    """random crop from an image.

     Parameters
     ----------
     image : numpy.array
         An image with dimension of [row, col, channel] (default).
     annotations : list of list of 5 elements
         Annotations [[class_id, c_x, c_y, w, h], [class_id, c_x, c_y, w, h], ...].
     dst_size : tuple of 2 int
        For cropped height and width.

     Returns
     -------
         numpy.array
             A processed image
         list of list of 5 float elements
             A list of new annotations.
    """
    hrg, wrg = dst_size
    image, classes, coords = tl.prepro.obj_box_crop(
        im=image,
        classes=[e[0] for e in annotations],
        coords=[e[1:] for e in annotations],
        wrg=wrg,
        hrg=hrg,
        is_rescale=True,
        is_center=True,
        is_random=True)
    annotations = [[a] + b for a, b in zip(classes, coords)]
    return image, annotations
